from django.test import TestCase
from django.contrib.auth import get_user_model
from .models import (
    AIProvider, AIModel, AIRole, AIConversation, 
    AIMessage, TokenUsageLog, MessageTemplate, UserFavorite
)

User = get_user_model()


class AIProviderModelTest(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username='testuser',
            email='test@example.com',
            password='testpass123'
        )
    
    def test_create_ai_provider(self):
        """测试创建AI提供商"""
        provider = AIProvider.objects.create(
            owner=self.user,
            name='Test OpenAI',
            provider_type='openai',
            api_endpoint='https://api.openai.com/v1',
            api_key='test-api-key',
            is_active=True
        )
        
        self.assertEqual(provider.name, 'Test OpenAI')
        self.assertEqual(provider.provider_type, 'openai')
        self.assertEqual(provider.owner, self.user)
        self.assertTrue(provider.is_active)
        
        # 测试权限检查方法
        self.assertTrue(provider.can_edit(self.user))
        self.assertTrue(provider.can_delete(self.user))
    
    def test_api_key_encryption(self):
        """测试API密钥加密"""
        original_key = 'test-api-key-123'
        provider = AIProvider.objects.create(
            owner=self.user,
            name='Test Provider',
            provider_type='openai',
            api_endpoint='https://api.openai.com/v1',
            api_key=original_key
        )
        
        # API密钥应该被加密存储
        self.assertNotEqual(provider.api_key, original_key)
        
        # 解密后应该能获取原始密钥
        decrypted_key = provider.decrypt_api_key()
        self.assertEqual(decrypted_key, original_key)


class AIModelModelTest(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username='testuser',
            email='test@example.com',
            password='testpass123'
        )
        self.provider = AIProvider.objects.create(
            owner=self.user,
            name='Test Provider',
            provider_type='openai',
            api_endpoint='https://api.openai.com/v1',
            api_key='test-key'
        )
    
    def test_create_ai_model(self):
        """测试创建AI模型"""
        model = AIModel.objects.create(
            provider=self.provider,
            name='GPT-4',
            model_code='gpt-4',
            max_tokens=8192,
            cost_per_1k_tokens=0.03
        )
        
        self.assertEqual(model.name, 'GPT-4')
        self.assertEqual(model.model_code, 'gpt-4')
        self.assertEqual(model.provider, self.provider)
        self.assertEqual(model.max_tokens, 8192)


class AIRoleModelTest(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username='testuser',
            email='test@example.com',
            password='testpass123'
        )
    
    def test_create_ai_role(self):
        """测试创建AI角色"""
        role = AIRole.objects.create(
            created_by=self.user,
            name='编程助手',
            system_prompt='你是一个专业的编程助手，帮助用户解决编程问题。',
            category='programming',
            is_public=False
        )
        
        self.assertEqual(role.name, '编程助手')
        self.assertEqual(role.category, 'programming')
        self.assertEqual(role.created_by, self.user)
        self.assertFalse(role.is_public)
        self.assertEqual(role.usage_count, 0)


class AIConversationModelTest(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username='testuser',
            email='test@example.com',
            password='testpass123'
        )
        self.provider = AIProvider.objects.create(
            owner=self.user,
            name='Test Provider',
            provider_type='openai',
            api_endpoint='https://api.openai.com/v1',
            api_key='test-key'
        )
        self.model = AIModel.objects.create(
            provider=self.provider,
            name='GPT-4',
            model_code='gpt-4'
        )
        self.role = AIRole.objects.create(
            created_by=self.user,
            name='助手',
            system_prompt='你是一个有用的助手。',
            category='assistant'
        )
    
    def test_create_conversation(self):
        """测试创建对话会话"""
        conversation = AIConversation.objects.create(
            owner=self.user,
            ai_model=self.model,
            ai_role=self.role,
            title='测试对话',
            status='active'
        )
        
        self.assertEqual(conversation.title, '测试对话')
        self.assertEqual(conversation.owner, self.user)
        self.assertEqual(conversation.ai_model, self.model)
        self.assertEqual(conversation.ai_role, self.role)
        self.assertEqual(conversation.status, 'active')
        self.assertEqual(conversation.used_tokens, 0)


class AIMessageModelTest(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username='testuser',
            email='test@example.com',
            password='testpass123'
        )
        self.provider = AIProvider.objects.create(
            owner=self.user,
            name='Test Provider',
            provider_type='openai',
            api_endpoint='https://api.openai.com/v1',
            api_key='test-key'
        )
        self.model = AIModel.objects.create(
            provider=self.provider,
            name='GPT-4',
            model_code='gpt-4'
        )
        self.role = AIRole.objects.create(
            created_by=self.user,
            name='助手',
            system_prompt='你是一个有用的助手。',
            category='assistant'
        )
        self.conversation = AIConversation.objects.create(
            owner=self.user,
            ai_model=self.model,
            ai_role=self.role,
            title='测试对话'
        )
    
    def test_create_user_message(self):
        """测试创建用户消息"""
        message = AIMessage.objects.create(
            conversation=self.conversation,
            content='你好，请帮我解答一个问题。',
            token_count=10
        )
        
        self.assertEqual(message.conversation, self.conversation)
        self.assertEqual(message.content, '你好，请帮我解答一个问题。')
        self.assertIsNone(message.ai_role)  # 用户消息没有AI角色
        self.assertTrue(message.is_user_message)
        self.assertEqual(message.token_count, 10)
    
    def test_create_ai_message(self):
        """测试创建AI消息"""
        message = AIMessage.objects.create(
            conversation=self.conversation,
            ai_role=self.role,
            content='你好！我很乐意帮助你解答问题。',
            token_count=15
        )
        
        self.assertEqual(message.conversation, self.conversation)
        self.assertEqual(message.ai_role, self.role)
        self.assertFalse(message.is_user_message)
        self.assertEqual(message.token_count, 15)


class TokenUsageLogModelTest(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(
            username='testuser',
            email='test@example.com',
            password='testpass123'
        )
        self.provider = AIProvider.objects.create(
            owner=self.user,
            name='Test Provider',
            provider_type='openai',
            api_endpoint='https://api.openai.com/v1',
            api_key='test-key'
        )
        self.model = AIModel.objects.create(
            provider=self.provider,
            name='GPT-4',
            model_code='gpt-4'
        )
        self.conversation = AIConversation.objects.create(
            owner=self.user,
            ai_model=self.model,
            title='测试对话'
        )
        self.message = AIMessage.objects.create(
            conversation=self.conversation,
            content='测试消息',
            token_count=20
        )
    
    def test_create_token_usage_log(self):
        """测试创建Token使用记录"""
        log = TokenUsageLog.objects.create(
            user=self.user,
            conversation=self.conversation,
            message=self.message,
            ai_model=self.model,
            input_tokens=10,
            output_tokens=15
        )
        
        self.assertEqual(log.user, self.user)
        self.assertEqual(log.input_tokens, 10)
        self.assertEqual(log.output_tokens, 15)
        self.assertEqual(log.total_tokens, 25)  # 自动计算总数
